﻿/*
* Copyright (C) 2020 ~ 2024 chaigec All Rights Reserved.
*
* Author:     chaigec <chaigec@163.com>
* Maintainer: chaigec <chaigec@163.com>

 * Use of this source code is governed by MIT license that can be found in the
 * LICENSE file in the root of the source tree. All contributing project authors
 * may be found in the AUTHORS file in the root of the source tree.
*/

#include "MkPsDemuxer.h"
#include "MkPsUtil.h"
#include "MkUtil/MkLog.h"

CMkPsDemuxer::CMkPsDemuxer()
    : m_LastVideoPts(0)
    , m_pLastEsBegin(nullptr)
    , m_nLastEsLen(0)
    , m_DecodeAacAdts(FALSE)
{
}

Uint32 CMkPsDemuxer::Input(Uint8 *pBuf, Uint32 nLen, Uint32& nUsed, BOOL bEnd/* = FALSE*/)
{
    Uint8 *pOffset = pBuf;
    Uint8 *secondHeader = nullptr;
    Uint32 ErrorCode = GetPsHeader(pBuf, nLen, &pOffset);
    MkCheckErrorLog(ErrorCode, "get first ps header failed\n");

    while (pOffset) {
        ErrorCode = GetPsHeader(pOffset + 4, static_cast<Uint32>(nLen - (pOffset - pBuf) - 4), &secondHeader);
        if (NoneError == ErrorCode) {
            InputPsPacket(pOffset, static_cast<Uint32>(secondHeader - pOffset));
        } else {
            break;
        }
        pOffset = secondHeader;
    }
    nUsed = pOffset - pBuf;
    if (bEnd) {
        InputPsPacket(pOffset, nLen - nUsed);
        nUsed = nLen;
    }
    return NoneError;
}

Uint32 CMkPsDemuxer::InputPsPacket(Uint8 *pPs, Uint32 nLen)
{
    Uint8* pSubHeader = nullptr;
    Uint8* pOffset = pPs;
    Uint8* pEnd = pPs + nLen;
    Uint32 ErrorCode = GetPsSubItem(pOffset, nLen, &pSubHeader);
    MkCheckErrorLog(ErrorCode, "get sub header failed\n");

    nLen -= (pSubHeader - pOffset);
    pOffset = pSubHeader;
    Uint32 PacketStartCode = 0;
    Uint32 nUsed = 0;
    while (pOffset < pEnd) {
        PacketStartCode = pOffset[0] << 24 | pOffset[1] << 16 | pOffset[2] << 8 | pOffset[3];
        if (PacketStartCode < 0x00000100 || PacketStartCode>0x000001FF) {
            MkWarningLog("error packet start code\n");
            pSubHeader = nullptr;
            ErrorCode = GetPsSubItem(pOffset, nLen, &pSubHeader);
            MkCheckErrorLog(ErrorCode, "get sub header failed\n");
            pOffset = pSubHeader;
            PacketStartCode = pOffset[0] << 24 | pOffset[1] << 16 | pOffset[2] << 8 | pOffset[3];
        }
        nUsed = 0;
        switch (PacketStartCode) {
        case PsStartCode:
            ErrorCode = ParsePsHeader(pOffset, nLen, nUsed);
            break;
        case PsSystemStartCode:
            ErrorCode = ParseSystemHeader(pOffset, nLen, nUsed);
            break;
        case PsMapStartCode:
            ErrorCode = ParsePsMap(pOffset, nLen, nUsed);
            break;
        default:
            ErrorCode = ParsePes(pOffset, nLen, PacketStartCode, nUsed);
            break;
        }
        MkCheckErrorLog(ErrorCode, "parse ps sub failed\n");
        pOffset += nUsed;
    }
    if (m_pLastEsBegin) {
        m_VideoPacket.BufferList.Append(m_pLastEsBegin, m_nLastEsLen);
        m_EsQueue.push(m_VideoPacket);
        m_VideoPacket.BufferList.Clear();
        m_VideoPacket.bKeyPacket = FALSE;
        m_VideoPacket.CodecType = MkCodecUnknown;
        m_VideoPacket.Pts = m_VideoPacket.Dts = 0;
        m_pLastEsBegin = nullptr;
        m_nLastEsLen = 0;
    }
    return NoneError;

    //Uint32 nUsed = 0;
    //Uint32 psLen = nLen;
    //Uint32 ErrorCode = ParsePsHeader(pPs, nLen, nUsed);
    //MkCheckErrorLog(ErrorCode, "parse ps header failed\n");
    //Uint8* pOffset = pPs + nUsed;
    //nLen -= nUsed;
    //nUsed = 0;
    //Uint32 PacketStartCode = 0;
    //while (nLen > nUsed) {
    //    nLen -= nUsed;
    //    pOffset += nUsed;
    //    nUsed = 0;
    //    PacketStartCode = pOffset[0] << 24 | pOffset[1] << 16 | pOffset[2] << 8 | pOffset[3];
    //    if (PacketStartCode < 0x00000100 || PacketStartCode>0x000001FF) {
    //        MkErrorLog("error packet start code\n");
    //        return NoneError;
    //    }
    //    switch (PacketStartCode) {
    //    case PsSystemStartCode:
    //        ErrorCode = ParseSystemHeader(pOffset, nLen, nUsed);
    //        break;
    //    case PsMapStartCode:
    //        ErrorCode = ParsePsMap(pOffset, nLen, nUsed);
    //        break;
    //    default:
    //        ErrorCode = ParsePes(pOffset, nLen, PacketStartCode, nUsed);
    //        break;
    //    }
    //    if (NoneError != ErrorCode) {
    //        MkCheckErrorLog(ErrorCode, "parse ps item failed\n");
    //    }
    //}
    //if (m_pLastEsBegin) {
    //    m_VideoPacket.BufferList.Append(m_pLastEsBegin, m_nLastEsLen);
    //    m_EsQueue.push(m_VideoPacket);
    //    m_VideoPacket.BufferList.Clear();
    //    m_VideoPacket.bKeyPacket = FALSE;
    //    m_VideoPacket.CodecType = MkCodecUnknown;
    //    m_VideoPacket.Pts = m_VideoPacket.Dts = 0;
    //    m_pLastEsBegin = nullptr;
    //    m_nLastEsLen = 0;
    //}
    //return NoneError;
}

Uint32 CMkPsDemuxer::GetEsPacket(MkEsPacket &Packet)
{
    if (m_EsQueue.empty()) {
        return QueueEmpty;
    }
    Packet = m_EsQueue.front();
    m_EsQueue.pop();
    return NoneError;
}


Uint32 CMkPsDemuxer::ParsePsHeader(Uint8* pBuf, Uint32 nLen, Uint32& nUsed)
{
    if (nLen < PsHeaderLen) {
        return InvalidParams;
    }
    Uint32 PacketStartCode = pBuf[0] << 24 | pBuf[1] << 16 | pBuf[2] << 8 | pBuf[3];
    if (PsStartCode != PacketStartCode) {
        return InvalidParams;
    }
    //MkPsHeader* psHeader = (MkPsHeader*)pBuf;
    //33 bit
    Uint64 SystemClockReference = (Uint64)(((pBuf[4] >> 3) & 0x07)) << 30;  //3bit
    SystemClockReference |= (pBuf[4] & 0x03) << 28;               //2bit
    SystemClockReference |= pBuf[5] << 20;                        //8bit
    SystemClockReference |= ((pBuf[6] >> 3) & 0x1F) << 15;        //5bit;
    SystemClockReference |= (pBuf[6] & 0x03) << 13;               //2bit;
    SystemClockReference |= pBuf[7] << 5;                         //8bit
    SystemClockReference |= (pBuf[8] >> 3) & 0x1f;                //5bit
    //9bit
    Uint32 SystemClockReferenceExtension = pBuf[8] & 0x03;
    SystemClockReferenceExtension |= (pBuf[9] >> 1) & 0x7F;

    Uint32 ProgramMuxRate = pBuf[10] << 14;
    ProgramMuxRate |= pBuf[11] << 6;
    ProgramMuxRate |= (pBuf[12] >> 2) & 0x3F;
    m_VideoParam.SetBitRate(ProgramMuxRate);
    Uint32 PackStuffingLenght = pBuf[13] & 0x07;
    if (PackStuffingLenght + PsHeaderLen > nLen) {
        return InvalidParams;
    }
    nUsed = PackStuffingLenght + PsHeaderLen;
    return NoneError;
}

Uint32 CMkPsDemuxer::ParseSystemHeader(Uint8* pBuf, Uint32 nLen, Uint32& nUsed)
{
    if (nLen < PsSystemHeaderLen) {
        return InvalidParams;
    }
    Uint32 HeaderLen = pBuf[4] << 8 | pBuf[5];
    if (nLen < 6 + HeaderLen) {
        return InvalidParams;
    }
    Uint32 RateBound = (pBuf[6] & 0x7F) << 15;
    RateBound |= pBuf[7] << 7;
    RateBound |= (pBuf[8] >> 1) & 0x7F;
    Uint8 AudioBound = pBuf[9] >> 2;
    Uint8 FixedFlag = (pBuf[9] >> 1) & 0x01;
    Uint8 CspsFlag = pBuf[9] & 0x01;
    Uint8 SystemAudioLockFlag = (pBuf[10] >> 7) & 0x01;
    Uint8 SystemVideoLockFlag = (pBuf[10] >> 6) & 0x01;
    Uint8  VideoBound = pBuf[10] & 0x1F;
    Uint8 PacketRateRestrictionFlag = (pBuf[11] >> 7) & 0x01;
    Uint8 *pOffset = pBuf + PsSystemHeaderLen;
    Uint8 pNextBit = (pOffset[0] >> 7) & 0x01;
    while (pNextBit && pOffset < pBuf + HeaderLen + 6) {
        Uint32 StreamId = pOffset[0];
        Uint8 Flag = pOffset[1] >> 6;
        Uint8 PstdBufferBoundScale = (pOffset[1] >> 5) & 0x01;
        Uint32 PstdBufferSizeBound = (pOffset[1] & 0x1F) << 8;
        PstdBufferSizeBound |= pOffset[2];
        Uint32 StartCode = (PsStartCode & 0xFFFFFF00) | StreamId;
        m_MapPsEsType[StartCode] = 0;
        pOffset += 3;
        pNextBit = (pOffset[0] >> 7) & 0x01;
    }
    nUsed = HeaderLen + 6;
    return NoneError;
}

Uint32 CMkPsDemuxer::ParsePsMap(Uint8* pBuf, Uint32 nLen, Uint32& nUsed)
{
    if (nLen < PsMapHeaderLen) {
        return InvalidParams;
    }
    Uint32 PsMapLen = pBuf[4] << 8 | pBuf[5];
    if (PsMapLen + 6 > nLen) {
        return InvalidParams;
    }
    nUsed = 6 + PsMapLen;
    Uint32 PsInfoLen = pBuf[8] << 8 | pBuf[9];
    Uint8 *pOffset = pBuf + PsMapHeaderLen;
    pOffset += PsInfoLen;
    Uint32 ElementaryStreamMapLenght = pOffset[0] << 8 | pOffset[1];
    pOffset += 2;
    while (pOffset + 4 < pBuf + nUsed) {
        Uint8 StreamType = pOffset[0];
        Uint8 StreamId = pOffset[1];
        Uint32 ElementaryStreamInfoLen = pOffset[2] << 8 | pOffset[3];
        Uint32 StartCode = (PsStartCode & 0xFFFFFF00) | StreamId;
        m_MapPsEsType[StartCode] = StreamType;
        MkCodecType Codec = CMkPsUtil::PsCodecToMkCodecType((MkPsStreamType)StreamType);
        if (MkCodecUnknown != Codec) {
            if (Codec > MkCodecCavs) {
                m_AudioParam.SetCodec(Codec);
                //ps stream payload g711a and g711u alway sample rate 8000HZ and channel 1
                if (MkCodecG711a == Codec
                    || MkCodecG711u == Codec) {
                    m_AudioParam.SetSampletRate(MkSampleRate8000Hz);
                    m_AudioParam.SetChannel(1);
                    m_AudioParam.SetStereo(FALSE);
                }
            } else {
                m_VideoParam.SetCodec(Codec);
            }
        }
        pOffset += 4 + ElementaryStreamInfoLen;
    }
    Uint32 Crc = pOffset[0] << 24 | pOffset[1] << 16 | pOffset[2] << 8 | pOffset[3];
    return NoneError;
}

Uint32 CMkPsDemuxer::ParsePes(Uint8 *pBuf, Uint32 nLen, const Uint32& PacketStartCode, Uint32& nUsed)
{
    MkMap<Uint32, Uint8>::iterator it = m_MapPsEsType.find(PacketStartCode);
    Uint32 PesLen = pBuf[4] << 8 | pBuf[5];
    if (PesLen + 6 > nLen) {  //6 bytes is start bytes 0x000001e0 + the pes len 2bytes
        return InvalidParams;
    }
    MkCodecType CodecType = MkCodecUnknown;
    Uint8* pEsBegin = nullptr;
    Uint32 nEsLen = 0;
    Uint64 Pts = 0;
    Uint64 Dts = 0;
    Uint8 *pOffset = pBuf + 6;
    if (it != m_MapPsEsType.end()
        && MkCodecUnknown != (CodecType = CMkPsUtil::PsCodecToMkCodecType((MkPsStreamType)it->second))) {
        Uint8 PtsDtsFlag = (pOffset[1] >> 6) & 0x03;
        Uint8 EsCrFlag = (pOffset[1] >> 5) & 0x01;
        Uint8 EsRateFlag = (pOffset[1] >> 4) & 0x01;
        Uint8 DsmTrickModeFlag = (pOffset[1] >> 3) & 0x01;
        Uint8 AdditionalCopyInfoFlag = (pOffset[1] >> 2) & 0x01;
        Uint8 PesCrcFlag = (pOffset[1] >> 1) & 0x01;
        Uint8 PesExtensionFlag = pOffset[1] & 0x01;
        Uint32 PesHeaderLen = pOffset[2];
        if (PesHeaderLen + 3 > PesLen) {  //3 bytes is upper parsed 3bytes
            return InvalidParams;
        }

        pOffset += 3;
        pEsBegin = pOffset + PesHeaderLen;
        nEsLen = PesLen - PesHeaderLen - 3;
        //only pts
        if (0x02 == PtsDtsFlag) {
            if (0x02 != ((pOffset[0] >> 4) & 0x0F)) {
                return InvalidParams;
            }
            Pts = ((Uint64)((pOffset[0] >> 1) & 0x07)) << 30;
            Pts |= pOffset[1] << 22;
            Pts |= ((pOffset[2] >> 1) & 0x7F) << 15;
            Pts |= pOffset[3] << 7;
            Pts |= (pOffset[4] >> 1) & 0x7F;
            Dts = Pts;
            pOffset += 5;
        } else if (0x03 == PtsDtsFlag) {
            if (0x03 != ((pOffset[0] >> 4) & 0x0F)) {
                return InvalidParams;
            }
            Pts = ((Uint64)((pOffset[0] >> 1) & 0x07)) << 30;
            Pts |= pOffset[1] << 22;
            Pts |= ((pOffset[2] >> 1) & 0x7F) << 15;
            Pts |= pOffset[3] << 7;
            Pts |= (pOffset[4] >> 1) & 0x7F;
            pOffset += 5;

            if (0x01 != ((pOffset[0] >> 4) & 0x0F)) {
                return InvalidParams;
            }
            Dts = ((Uint64)((pOffset[0] >> 1) & 0x07)) << 30;
            Dts |= pOffset[1] << 22;
            Dts |= ((pOffset[2] >> 1) & 0x7F) << 15;
            Dts |= pOffset[3] << 7;
            Dts |= (pOffset[4] >> 1) & 0x7F;
            pOffset += 5;
        }

        //video
        if (CodecType <= MkCodecCavs) {
            if (0 == m_VideoPacket.Dts) {
                m_VideoPacket.Dts = static_cast<Uint32>(Dts / (MkDefaultClockRate / 1000));
                m_VideoPacket.Pts = static_cast<Uint32>(Pts / (MkDefaultClockRate / 1000));
                m_VideoPacket.CodecType = CodecType;
            }
            if (0 == m_VideoParam.GetFrameRate() && 0 != m_LastVideoPts && m_LastVideoPts != m_VideoPacket.Pts) {
                m_VideoParam.SetFrameRate(1000 / (m_VideoPacket.Pts - m_LastVideoPts));
            }
            m_LastVideoPts = m_VideoPacket.Pts;
            if (0 == memcmp(&NaluStartCode, pEsBegin, sizeof(Uint32))) {
                pEsBegin += sizeof(Uint32);
                nEsLen -= sizeof(Uint32);
                Uint8 NaluType = 0;
                switch (CodecType) {
                case MkCodecH264:
                    CMkMediaUtil::Get264NaluType(*pEsBegin, NaluType);
                    if (MkH264NalSps == NaluType) {
                        m_VideoParam.SetSps(MkString((Int8*)pEsBegin, nEsLen));
                        Uint32 Width;
                        Uint32 Height;
                        Uint32 FrameRate;
                        CMkMediaUtil::H264SpsParse(pEsBegin, nEsLen, Width, Height, FrameRate);
                        m_VideoParam.SetWidth(Width);
                        m_VideoParam.SetHeight(Height);
                        m_VideoParam.SetFrameRate(FrameRate);
                        m_VideoPacket.BufferList.Append(pEsBegin, nEsLen);
                    } else if (MkH264NalPps == NaluType) {
                        m_VideoParam.SetPps(MkString((Int8 *)pEsBegin, nEsLen));
                        m_VideoPacket.BufferList.Append(pEsBegin, nEsLen);
                    } else {
                        if (MkH264NalIdrSlice == NaluType) {
                            m_VideoPacket.bKeyPacket = TRUE;
                        }
                        if (m_pLastEsBegin) {
                            m_VideoPacket.BufferList.Append(m_pLastEsBegin, m_nLastEsLen);
                        }
                        m_pLastEsBegin = pEsBegin;
                        m_nLastEsLen = nEsLen;
                    }
                    break;
                case MkCodecH265:
                    CMkMediaUtil::Get265NaluType(*pEsBegin, NaluType);
                    if (MkHevcNalSps == NaluType) {
                        m_VideoParam.SetSps(MkString((Int8 *)pEsBegin, nEsLen));
                        Uint32 Width;
                        Uint32 Height;
                        CMkMediaUtil::H265SpsParse(pEsBegin, nEsLen, Width, Height);
                        m_VideoParam.SetWidth(Width);
                        m_VideoParam.SetHeight(Height);
                        m_VideoPacket.BufferList.Append(pEsBegin, nEsLen);
                    } else if (MkHevcNalPps == NaluType) {
                        m_VideoParam.SetPps(MkString((Int8 *)pEsBegin, nEsLen));
                        m_VideoPacket.BufferList.Append(pEsBegin, nEsLen);
                    } else if (MkHevcNalVps == NaluType) {
                        m_VideoParam.SetVps(MkString((Int8 *)pEsBegin, nEsLen));
                        m_VideoPacket.BufferList.Append(pEsBegin, nEsLen);
                    } else {
                        if (MkHevcNalIdrNLp == NaluType || MkHevcNalIdrWRadl == NaluType) {
                            m_VideoPacket.bKeyPacket = TRUE;
                        }
                        if (m_pLastEsBegin) {
                            m_VideoPacket.BufferList.Append(m_pLastEsBegin, m_nLastEsLen);
                        }
                        m_pLastEsBegin = pEsBegin;
                        m_nLastEsLen = nEsLen;
                    }
                    break;
                default:
                    break;
                }
            } else {
                if (!m_pLastEsBegin) {
                    return InvalidParams;
                }
                memmove(m_pLastEsBegin + m_nLastEsLen, pEsBegin, nEsLen);
                m_nLastEsLen += nEsLen;
            }
        } else {  //audio
            if (m_pLastEsBegin) {
                m_VideoPacket.BufferList.Append(m_pLastEsBegin, m_nLastEsLen);
                m_EsQueue.push(m_VideoPacket);
                m_VideoPacket.BufferList.Clear();
                m_VideoPacket.bKeyPacket = FALSE;
                m_VideoPacket.CodecType = MkCodecUnknown;
                m_VideoPacket.Pts = m_VideoPacket.Dts = 0;
                m_pLastEsBegin = nullptr;
                m_nLastEsLen = 0;
            }
            MkEsPacket  AudioPacket;
            AudioPacket.Dts = static_cast<Uint32>(Dts / (MkDefaultClockRate / 1000));
            AudioPacket.Pts = static_cast<Uint32>(Pts / (MkDefaultClockRate / 1000));
            AudioPacket.CodecType = CodecType;

            if (MkCodecAac == CodecType) {
                if (!m_DecodeAacAdts) {
                    MkAacObjectType ObjectType;
                    MkAudioSampleRate SampleRate;
                    Uint8 nChannel;
                    CMkMediaUtil::AacAdtsParse(pEsBegin, nEsLen, ObjectType, SampleRate, nChannel);
                    m_AudioParam.SetObjectType(ObjectType);
                    m_AudioParam.SetChannel(nChannel);
                    m_AudioParam.SetSampletRate(SampleRate);
                    m_DecodeAacAdts = TRUE;
                }
                AudioPacket.BufferList.Append(pEsBegin + MkAacAdtsLen, nEsLen - MkAacAdtsLen);
            } else {
                AudioPacket.BufferList.Append(pEsBegin, nEsLen);
            }
            m_EsQueue.push(AudioPacket);
        }
    }
    nUsed = 6 + PesLen;
    return NoneError;
}

Uint32 CMkPsDemuxer::GetPsHeader(Uint8 *pBuf, Uint32 nLen, Uint8 **psHeader)
{
    Uint8 *pOffset = pBuf;
    Uint8 *pEnd = pBuf + nLen;
    while (pOffset < pEnd) {
        if (0x00 == pOffset[0]
            && 0x00 == pOffset[1]) {
            if (0x00 == pOffset[2]) {
                pOffset += 3;
                continue;
            } else if (0x01 == pOffset[2]
                && 0xba == pOffset[3]) {
                *psHeader = pOffset;
                return NoneError;
            }
        }
        pOffset++;
    }
    return SystemError;
}

Uint32 CMkPsDemuxer::GetPsSubItem(Uint8 *pBuf, Uint32 nLen, Uint8 **ppHeader)
{
    Uint8 *pOffset = pBuf;
    Uint8 *pEnd = pBuf + nLen;
    while (pOffset < pEnd) {
        if (0x00 == pOffset[0]
            && 0x00 == pOffset[1]
            && 0x01 == pOffset[2]
            && 0xA0 < pOffset[3]) {
            *ppHeader = pOffset;
            return NoneError;
        }
        pOffset++;
    }
    return NotHasEnoughBuffer;
}

CMkPsDemuxer::~CMkPsDemuxer()
{
    while (!m_EsQueue.empty()) {
        m_EsQueue.pop();
    }
}
